package rtsp_repush

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"os/exec"
	"strings"

	"github.com/bluenviron/gortsplib/v3"
	"github.com/bluenviron/gortsplib/v3/pkg/formats"
	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtph264"
	"github.com/bluenviron/gortsplib/v3/pkg/media"
	"github.com/bluenviron/gortsplib/v3/pkg/url"
	"github.com/bluenviron/mediacommon/pkg/codecs/h264"
	"github.com/pion/rtp"
)

// RePush IDE帧前面添加SPS/PPS
// "rtsp://127.0.0.1:554/live/movie"
// "rtsp://127.0.0.1:554/live/movie2"
// ffmpeg -re -i rtmp://127.0.0.1/live/movie -c:v libx264 -an -bf 0 -f rtsp rtsp://127.0.0.1/live/movie
func RePush(ctx context.Context, inputURL string, outputURL string) error {
	c := gortsplib.Client{}

	// parse URL
	u, err := url.Parse(inputURL)
	if err != nil {
		log.Println(err)
		return err
	}

	// connect to the server
	err = c.Start(u.Scheme, u.Host)
	if err != nil {
		log.Println(err)
		return err
	}
	defer c.Close()

	// find published medias
	medias, baseURL, _, err := c.Describe(u)
	if err != nil {
		log.Println(err)
		return err
	}

	// find the H264 media and format
	var forma *formats.H264
	medi := medias.FindFormat(&forma)
	if medi == nil {
		log.Println(err)
		return err
	}

	// setup RTP/H264 -> H264 decoder
	rtpDec, err := forma.CreateDecoder2()
	if err != nil {
		log.Println(err)
		return err
	}
	rtpEnc, _ := forma.CreateEncoder2()

	// setup a single media
	_, err = c.Setup(medi, baseURL, 0, 0)
	if err != nil {
		log.Println(err)
		return err
	}

	var targetCli = &gortsplib.Client{}
	if err = targetCli.StartRecording(outputURL, media.Medias{medi}); err != nil {
		log.Println(err)
		return err
	}
	defer c.Close()

	// called when a RTP packet arrives
	c.OnPacketRTP(medi, forma, func(pkt *rtp.Packet) {
		// extract access units from RTP packets
		au, pts, err := rtpDec.DecodeUntilMarker(pkt)
		if err != nil {
			if err != rtph264.ErrNonStartingPacketAndNoPrevious && err != rtph264.ErrMorePacketsNeeded {
				log.Printf("ERR: %v", err)
			}
			return
		}
		_ = pts

		for _, nalu := range au {
			t := h264NALUType(nalu[0])
			if t == h264.NALUTypeIDR {
				{
					packets, err := rtpEnc.Encode([][]byte{forma.SPS}, pts)
					if err != nil {
						log.Println(err)
						return
					}

					for _, packet := range packets {
						_ = packet
						if err = targetCli.WritePacketRTP(medi, packet); err != nil {
							log.Println(err)
						}
					}
				}
				{
					packets, err := rtpEnc.Encode([][]byte{forma.PPS}, pts)
					if err != nil {
						log.Println(err)
						return
					}

					for _, packet := range packets {
						_ = packet
						if err = targetCli.WritePacketRTP(medi, packet); err != nil {
							log.Println(err)
						}
					}
				}
			}

			switch t {
			case h264.NALUTypeIDR:
				log.Println(t)
			case h264.NALUTypeNonIDR:
			default:
				log.Println(t)
			}
		}

		packets, err := rtpEnc.Encode(au, pts)
		if err != nil {
			log.Println(err)
			return
		}

		for _, packet := range packets {
			_ = packet
			if err = targetCli.WritePacketRTP(medi, packet); err != nil {
				log.Println(err)
			}
		}
	})

	// start playing
	_, err = c.Play(nil)
	if err != nil {
		log.Println(err)
		return err
	}

	go func() {
		<-ctx.Done()
		_ = c.Close()
	}()

	// wait until a fatal error
	_ = c.Wait()

	return nil
}

func h264NALUType(firstByte byte) h264.NALUType {
	t := h264.NALUType((firstByte & 0x1F) >> 0)
	return t
}

// RtmpToRtsp ffmpeg: rtmp流转换为rtsp流, 重新编码, 去掉B帧
func RtmpToRtsp(ctx context.Context, inputURI, outputURI string) error {
	var rr = func(head string, rd io.Reader) {
		r := bufio.NewReader(rd)
		for {
			l, _, err := r.ReadLine()
			if err != nil {
				return
			}
			log.Printf("[%v] %v", head, string(l))
		}
	}

	strCmd := fmt.Sprintf("ffmpeg -re -i %v -c:v libx264 -an -bf 0 -f rtsp %v", inputURI, outputURI)
	log.Printf("%v", strCmd)

	args := strings.Split(strCmd, " ")
	cmd := exec.CommandContext(ctx, args[0], args[1:]...)

	errR, _ := cmd.StderrPipe()
	stdR, _ := cmd.StdoutPipe()

	go rr("stderr", errR)
	go rr("stdout", stdR)

	if err := cmd.Run(); err != nil {
		log.Println(err)
	} else {
		_ = cmd.Wait()
	}
	return nil
}
